#include "kernel_operator.h"

using namespace AscendC;

constexpr uint32_t BUFFER_NUM = 2;

template <typename T>
class KernelLogSumExp
{
    using Tfp32 = float;

public:
    __aicore__ inline KernelLogSumExp() {}
    __aicore__ inline void Init(TPipe *pipeIn, GM_ADDR x, GM_ADDR y,
                                int32_t y_dimensional,
                                int32_t *y_ndarray, int32_t *x_ndarray,
                                int32_t *y_sumndarray, int32_t *x_sumndarray)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->y_dimensional = y_dimensional;

        this->y_ndarray = y_ndarray;
        this->x_ndarray = x_ndarray;
        this->y_sumndarray = y_sumndarray;
        this->x_sumndarray = x_sumndarray;

        xGm.SetGlobalBuffer((__gm__ T *)x, 1);
        yGm.SetGlobalBuffer((__gm__ T *)y, 1);
        pipe = pipeIn;

        pipe->InitBuffer(inQueueX, BUFFER_NUM, 32);
    }
    __aicore__ inline void Process()
    {
        LocalTensor<Tfp32> xLocal = inQueueX.AllocTensor<Tfp32>();
        Tfp32 xMax = 3;

        int dim = this->y_dimensional;
        for (int j = 0; j < this->y_sumndarray[dim]; j++)
        {
            yGm.SetValue(j, (T)0);
        }

        for (int j = 0; j < this->x_sumndarray[dim]; j++)
        {
            int y_start = 0;
            for (int k = 0; k < dim; k++)
            {
                if (this->y_ndarray[k] != 1)
                {
                    y_start += this->y_sumndarray[k] * (j / this->x_sumndarray[k] % this->x_ndarray[k]);
                }
            }
            Tfp32 x = static_cast<Tfp32>(xGm.GetValue(j));
            Tfp32 y = static_cast<Tfp32>(yGm.GetValue(y_start));

            xLocal.SetValue(0, x - xMax);
            Exp(xLocal, xLocal, 1);

            yGm.SetValue(y_start, static_cast<T>(xLocal.GetValue(0) + y));
        }
        for (int j = 0; j < this->y_sumndarray[dim]; j++)
        {
            Tfp32 y = static_cast<Tfp32>(yGm.GetValue(j));
            xLocal.SetValue(0, y);
            Ln(xLocal, xLocal, 1);
            yGm.SetValue(j, static_cast<T>(xMax + xLocal.GetValue(0)));
        }
        inQueueX.FreeTensor(xLocal);
    }

private:
    TPipe *pipe;
    TQue<QuePosition::VECIN, 1> inQueueX;

    GlobalTensor<T> xGm;
    GlobalTensor<T> yGm;

    int32_t y_dimensional;
    int32_t *y_ndarray;
    int32_t *x_ndarray;

    int32_t *y_sumndarray;
    int32_t *x_sumndarray;
};
extern "C" __global__ __aicore__ void log_sum_exp(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    TPipe pipe;
    GET_TILING_DATA(tiling_data, tiling);
    KernelLogSumExp<DTYPE_X> op;
    op.Init(&pipe, x, y,
            tiling_data.x_dimensional,
            tiling_data.y_ndarray, tiling_data.x_ndarray,
            tiling_data.y_sumndarray, tiling_data.x_sumndarray);
    op.Process();
}